// Fill out your copyright notice in the Description page of Project Settings.

#include "ScreenCapturer.h"
#include <Rendering/Texture2DResource.h>
#include <opencv2/imgproc.hpp>

AScreenCapturer::AScreenCapturer()
{
 	PrimaryActorTick.bCanEverTick = true;
	PrimaryActorTick.TickGroup = TG_PostUpdateWork;

    RootComponent = CreateDefaultSubobject<USceneComponent>(TEXT("RootComponent"));
    _captureCamera = CreateDefaultSubobject<UCameraComponent>(TEXT("CaptureCamera"));
    _captureCamera->SetupAttachment(RootComponent);

    _sceneCapture = CreateDefaultSubobject<USceneCaptureComponent2D>(TEXT("SceneCapturer"));
    _sceneCapture->SetupAttachment(_captureCamera);
}

void AScreenCapturer::BeginPlay()
{
    _writer.open(TCHAR_TO_UTF8(*_outputFilePath),
                 cv::VideoWriter::fourcc('M', 'J', 'P', 'G'), 30,
                 cv::Size(_resolution, _resolution), true);
    UE_LOG(LogTemp, Warning, TEXT("Video writer status: %d"), _writer.isOpened());

    _captureCamera->FieldOfView = _fieldOfView;
    _sceneCapture->FOVAngle = _fieldOfView;

    _renderTarget = NewObject<UTextureRenderTarget2D>();
    _renderTarget->InitCustomFormat(
        _resolution, _resolution, 
        EPixelFormat::PF_B8G8R8A8, true);
    _renderTarget->UpdateResourceImmediate();

    _sceneCapture->CaptureSource = SCS_FinalColorLDR;
    _sceneCapture->TextureTarget = _renderTarget;
    _sceneCapture->bCaptureEveryFrame = true;
    _sceneCapture->bAlwaysPersistRenderingState = true;
	Super::BeginPlay();
}

void AScreenCapturer::EndPlay(const EEndPlayReason::Type reason)
{
    _writer.release();
    Super::EndPlay(reason);
}

void AScreenCapturer::Tick(float DeltaTime)
{
    auto resource =
        _renderTarget->GameThread_GetRenderTargetResource();
    if (resource)
    {
        TArray<FColor> buffer;
        resource->ReadPixels(buffer);

        cv::Mat frame(_renderTarget->GetSurfaceHeight(), _renderTarget->GetSurfaceWidth(),
                      CV_8UC4, buffer.GetData());
        cv::cvtColor(frame, frame, cv::COLOR_RGBA2RGB);
        _writer.write(frame);
    }
	Super::Tick(DeltaTime);
}
